{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Robotic Manipulation - Bin Picking.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.5-final"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TKvYiJgnYExi"
      },
      "source": [
        "**I recommend you run the first code cell of this notebook immediately, to start provisioning drake on the cloud machine, then you can leave this window open as you [read the textbook](http://manipulation.csail.mit.edu/clutter.html).**\n",
        "\n",
        "# Notebook Setup\n",
        "\n",
        "The following cell will:\n",
        "- on Colab (only), install Drake to `/opt/drake`, install Drake's prerequisites via `apt`, and add pydrake to `sys.path`.  This will take approximately two minutes on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  If you navigate between notebooks using Colab's \"File->Open\" menu, then you can avoid provisioning a separate machine for each notebook.\n",
        "- define some utility methods/classes that will eventually disappear from this notebook and live in drake.\n",
        "\n",
        "You will need to rerun this cell if you restart the kernel, but it should be fast because the machine will already have drake installed."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A4QOaw_zYLfI"
      },
      "source": [
        "import importlib\n",
        "import sys\n",
        "from urllib.request import urlretrieve\n",
        "\n",
        "if 'google.colab' in sys.modules and importlib.util.find_spec('manipulation') is None:\n",
        "    urlretrieve(f\"http://manipulation.csail.mit.edu/scripts/setup/setup_manipulation_colab.py\",\n",
        "                \"setup_manipulation_colab.py\")\n",
        "    from setup_manipulation_colab import setup_manipulation\n",
        "    setup_manipulation(manipulation_sha='master', drake_version='latest', drake_build='continuous')\n",
        "\n",
        "# Determine if this notebook is currently running as a notebook or a unit test.\n",
        "from IPython import get_ipython\n",
        "running_as_notebook = get_ipython() and hasattr(get_ipython(), 'kernel')\n",
        "\n",
        "# Setup rendering (with xvfb), if necessary:\n",
        "import os\n",
        "if 'google.colab' in sys.modules and os.getenv(\"DISPLAY\") is None:\n",
        "    from pyvirtualdisplay import Display\n",
        "    display = Display(visible=0, size=(1400, 900))\n",
        "    display.start()\n",
        "\n",
        "server_args = []\n",
        "if 'google.colab' in sys.modules:\n",
        "  server_args = ['--ngrok_http_tunnel']\n",
        "\n",
        "# Start two meshcat servers (one for 2d) to use for the remainder of this notebook.\n",
        "from meshcat.servers.zmqserver import start_zmq_server_as_subprocess\n",
        "proc_2d, zmq_url_2d, web_url_2d = start_zmq_server_as_subprocess(server_args=server_args)\n",
        "proc, zmq_url, web_url = start_zmq_server_as_subprocess(server_args=server_args)\n",
        "\n",
        "# Imports\n",
        "import numpy as np\n",
        "from IPython.display import display, HTML\n",
        "from ipywidgets import Textarea\n",
        "\n",
        "from pydrake.all import ( \n",
        "    AddMultibodyPlantSceneGraph, ConnectMeshcatVisualizer, \n",
        "    DiagramBuilder, RigidTransform, RotationMatrix, Box,    \n",
        "    CoulombFriction, FindResourceOrThrow, FixedOffsetFrame, \n",
        "    GeometryInstance, MeshcatContactVisualizer, Parser, PlanarJoint,  \n",
        "    JointIndex, Simulator, ProcessModelDirectives, LoadModelDirectives\n",
        ")\n",
        "\n",
        "from functools import partial\n",
        "import open3d as o3d\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import display, HTML\n",
        "import meshcat\n",
        "import meshcat.geometry as g\n",
        "import meshcat.transformations as tf\n",
        "\n",
        "from pydrake.all import (\n",
        "    ConnectPlanarSceneGraphVisualizer,\n",
        "    ConnectDrakeVisualizer, DepthCameraProperties, RgbdSensor,\n",
        "    RandomGenerator, UniformlyRandomRotationMatrix, RollPitchYaw,\n",
        "    MakeRenderEngineVtk, RenderEngineVtkParams, Role, UnitInertia, set_log_level\n",
        ")\n",
        "\n",
        "from ipywidgets import Dropdown, FloatSlider, Layout\n",
        "from pydrake.all import (\n",
        "    Sphere, Cylinder, Box, Capsule, Ellipsoid, SpatialInertia)\n",
        "\n",
        "\n",
        "#from pydrake.multibody.jupyter_widgets import MakeJointSlidersThatPublishOnCallback\n",
        "from manipulation.jupyter_widgets import MakeJointSlidersThatPublishOnCallback\n",
        "from manipulation.meshcat_utils import draw_open3d_point_cloud, draw_points\n",
        "from manipulation.open3d_utils import create_open3d_point_cloud\n",
        "from manipulation.mustard_depth_camera_example import MustardExampleSystem\n",
        "from manipulation.scenarios import AddRgbdSensors\n",
        "from manipulation.utils import FindResource\n",
        "\n",
        "set_log_level(\"warn\");\n",
        "\n",
        "ycb = [(\"cracker\", \"003_cracker_box.sdf\"), \n",
        "        (\"sugar\", \"004_sugar_box.sdf\"), \n",
        "        (\"soup\", \"005_tomato_soup_can.sdf\"), \n",
        "        (\"mustard\", \"006_mustard_bottle.sdf\"), \n",
        "        (\"gelatin\", \"009_gelatin_box.sdf\"), \n",
        "        (\"meat\", \"010_potted_meat_can.sdf\")]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fGbe-rJGJlF0"
      },
      "source": [
        "# Falling things (in 2D)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_JIDhLfN3BSe",
        "tags": []
      },
      "source": [
        "def clutter_gen():\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "\n",
        "    box = Box(10., 10., 10.)\n",
        "    X_WBox = RigidTransform([0, 0, -5])\n",
        "    mu = 1.0\n",
        "    plant.RegisterCollisionGeometry(plant.world_body(), X_WBox, box, \"ground\", CoulombFriction(mu, mu))\n",
        "    plant.RegisterVisualGeometry(plant.world_body(), X_WBox, box, \"ground\", [.9, .9, .9, 1.0])\n",
        "    planar_joint_frame = plant.AddFrame(FixedOffsetFrame(\"planar_joint_frame\", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
        "\n",
        "    parser = Parser(plant)\n",
        "\n",
        "    sdf = FindResourceOrThrow(\"drake/examples/manipulation_station/models/061_foam_brick.sdf\")\n",
        "\n",
        "    for i in range(20 if running_as_notebook else 2):\n",
        "        instance = parser.AddModelFromFile(sdf, f\"object{i}\")\n",
        "        plant.AddJoint(PlanarJoint(f\"joint{i}\", planar_joint_frame, plant.GetFrameByName(\"base_link\", instance), damping=[0,0,0]))\n",
        "\n",
        "    plant.Finalize()\n",
        "\n",
        "    vis = ConnectPlanarSceneGraphVisualizer(\n",
        "        builder, \n",
        "        scene_graph,\n",
        "        xlim=[-.6, .6],\n",
        "        ylim=[-.1, 0.5],\n",
        "        show=False,\n",
        "    )\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    simulator = Simulator(diagram)\n",
        "    plant_context = plant.GetMyContextFromRoot(simulator.get_mutable_context())\n",
        "\n",
        "    rs = np.random.RandomState()\n",
        "    z = 0.1\n",
        "    for i in range(plant.num_joints()):\n",
        "        joint = plant.get_joint(JointIndex(i))\n",
        "        joint.set_pose(plant_context, [rs.uniform(-.4, .4), z], rs.uniform(-np.pi/2.0, np.pi/2.0))\n",
        "        z += 0.1\n",
        "\n",
        "    vis.start_recording()\n",
        "    simulator.AdvanceTo(1.5 if running_as_notebook else 0.1)\n",
        "    vis.stop_recording()\n",
        "    ani = vis.get_recording_as_animation(repeat=False)\n",
        "    display(HTML(ani.to_jshtml()))\n",
        "    \n",
        "clutter_gen()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GoxjuO5PJlF4"
      },
      "source": [
        "# Falling things (in 3D)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K2oyAnHfJlF5"
      },
      "source": [
        "def clutter_gen():\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0005)\n",
        "\n",
        "    parser = Parser(plant)\n",
        "\n",
        "    parser.AddModelFromFile(FindResourceOrThrow(\n",
        "        \"drake/examples/manipulation_station/models/bin.sdf\"))\n",
        "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"bin_base\"))\n",
        "\n",
        "\n",
        "    rs = np.random.RandomState()  # this is for python\n",
        "    generator = RandomGenerator(rs.randint(1000))  # this is for c++\n",
        "    for i in range(10 if running_as_notebook else 2):\n",
        "        object_num = rs.randint(len(ycb))\n",
        "        sdf = FindResourceOrThrow(\"drake/manipulation/models/ycb/sdf/\" + ycb[object_num][1])\n",
        "        parser.AddModelFromFile(sdf, f\"object{i}\")\n",
        "\n",
        "    plant.Finalize()\n",
        "\n",
        "    renderer = \"my_renderer\"\n",
        "    scene_graph.AddRenderer(\n",
        "        renderer, MakeRenderEngineVtk(RenderEngineVtkParams()))\n",
        "    properties = DepthCameraProperties(width=640,\n",
        "                                       height=480,\n",
        "                                       fov_y=np.pi / 4.0,\n",
        "                                       renderer_name=renderer,\n",
        "                                       z_near=0.1,\n",
        "                                       z_far=10.0)\n",
        "    camera = builder.AddSystem(\n",
        "        RgbdSensor(parent_id=scene_graph.world_frame_id(),\n",
        "                   X_PB=RigidTransform(\n",
        "                       RollPitchYaw(np.pi, 0, np.pi/2.0),\n",
        "                       [0, 0, .8]),\n",
        "                   properties=properties,\n",
        "                   show_window=False))\n",
        "    camera.set_name(\"rgbd_sensor\")\n",
        "    builder.Connect(scene_graph.get_query_output_port(),\n",
        "                    camera.query_object_input_port())\n",
        "    builder.ExportOutput(camera.color_image_output_port(), \"color_image\")\n",
        "\n",
        "    # Note: if you're running this on a local machine, then you can \n",
        "    # use drake_visualizer to see the simulation.  (It's too slow to \n",
        "    # show the meshes on meshcat).\n",
        "    vis = ConnectDrakeVisualizer(\n",
        "        builder, \n",
        "        scene_graph\n",
        "    )\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    simulator = Simulator(diagram)\n",
        "    context = simulator.get_mutable_context()\n",
        "    plant_context = plant.GetMyContextFromRoot(context)\n",
        "\n",
        "    z = 0.1\n",
        "    for body_index in plant.GetFloatingBaseBodies():\n",
        "        tf = RigidTransform(\n",
        "                UniformlyRandomRotationMatrix(generator),  \n",
        "                [rs.uniform(-.15,.15), rs.uniform(-.2, .2), z])\n",
        "        plant.SetFreeBodyPose(plant_context, \n",
        "                              plant.get_body(body_index),\n",
        "                              tf)\n",
        "        z += 0.1\n",
        "\n",
        "    simulator.AdvanceTo(1.0 if running_as_notebook else 0.1)\n",
        "    color_image = diagram.GetOutputPort(\"color_image\").Eval(context)\n",
        "    plt.figure()\n",
        "    plt.imshow(color_image.data)\n",
        "    plt.axis('off')\n",
        "\n",
        "clutter_gen()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7aGJ9e3lJlF8"
      },
      "source": [
        "# Contact force \"inspector\"\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7dimkJhd9DjP",
        "tags": []
      },
      "source": [
        "def contact_force_inspector(slope=0.0, mu=1.0, second_brick=False):\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.01)\n",
        "\n",
        "    box = Box(10., 10., 10.)\n",
        "    X_WBox = RigidTransform(RotationMatrix.MakeYRotation(slope), [0, 0, -5.05])\n",
        "    plant.RegisterCollisionGeometry(plant.world_body(), X_WBox, box, \"ground\", CoulombFriction(mu, mu))\n",
        "    plant.RegisterVisualGeometry(plant.world_body(), X_WBox, box, \"ground\", [.9, .9, .9, 1.0])\n",
        "\n",
        "    parser = Parser(plant)\n",
        "    brick_sdf = FindResource(\"models/planar_foam_brick_collision_as_visual.sdf\")\n",
        "    parser.AddModelFromFile(brick_sdf)\n",
        "    frame = plant.AddFrame(FixedOffsetFrame(\"planar_joint_frame\", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
        "    plant.AddJoint(PlanarJoint(\"brick\", frame, plant.GetFrameByName(\"base_link\"), damping=[0,0,0]))\n",
        "\n",
        "    if second_brick:\n",
        "        brick2 = parser.AddModelFromFile(brick_sdf, \"brick2\")\n",
        "        plant.AddJoint(PlanarJoint(\"brick2\", frame, plant.GetFrameByName(\"base_link\", brick2), damping=[0,0,0]))\n",
        "\n",
        "    plant.Finalize()\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(\n",
        "        builder, \n",
        "        scene_graph, \n",
        "        zmq_url=zmq_url_2d,\n",
        "        server_args=server_args)\n",
        "    meshcat.set_planar_viewpoint(xmin=-.2, xmax=.2, ymin=-.2, ymax=0.3)\n",
        "\n",
        "    contact_visualizer = builder.AddSystem(MeshcatContactVisualizer(meshcat, \n",
        "                                           force_threshold=0.0,\n",
        "                                           contact_force_scale=1.0, \n",
        "                                           contact_force_radius=0.002,   \n",
        "                                           plant=plant))\n",
        "    builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
        "                    contact_visualizer.GetInputPort(\"pose_bundle\"))\n",
        "    builder.Connect(plant.get_contact_results_output_port(),\n",
        "                    contact_visualizer.GetInputPort(\"contact_results\"))\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    meshcat.load()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "\n",
        "    textoutput = Textarea(value=\"\", description=\"contact info: \", layout={'width':'800px','height':'100px'}, style={'description_width':'initial'})\n",
        "    display(textoutput)\n",
        "    def my_callback(plant_context):\n",
        "        results = plant.get_contact_results_output_port().Eval(plant_context)\n",
        "        textoutput.value = \"\" if results.num_point_pair_contacts()>0 else \"no contact\"\n",
        "        for i in range(results.num_point_pair_contacts()):\n",
        "            info = results.point_pair_contact_info(i)\n",
        "            pair = info.point_pair()\n",
        "            textoutput.value += f\"slip speed:{info.slip_speed():.4f}, depth:{pair.depth:.4f}, \"\n",
        "            textoutput.value += \"force:\" + np.array2string(info.contact_force(), formatter={'float': lambda x: \"{:5.2f}\".format(x)}) + \"\\n\"\n",
        "\n",
        "    lower_limit = [-0.1, -0.1, -np.pi/2.0]\n",
        "    upper_limit = [0.1, 0.1, np.pi/2.0]\n",
        "    if second_brick:\n",
        "        lower_limit += lower_limit\n",
        "        upper_limit += upper_limit\n",
        "        plant.SetPositions(\n",
        "            plant.GetMyContextFromRoot(context),\n",
        "            [0, 0, 0, 0.07, 0.07, 0.0])\n",
        "\n",
        "    MakeJointSlidersThatPublishOnCallback(\n",
        "        plant, diagram, context, my_callback, resolution=0.001,\n",
        "        lower_limit=lower_limit, upper_limit=upper_limit)    \n",
        "\n",
        "contact_force_inspector(\n",
        "    slope=0.1, \n",
        "    mu=0.5,\n",
        "    second_brick=True);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ctYY0SHJlGB"
      },
      "source": [
        "# Contact results \"inspector\"\n",
        "\n",
        "This simple visualization shows some of the complexity of the contact geometry problem.  I will make it better, but right now, when you move the objects into contact of each other you will see three points:  the contact point is in **red**, the contact normal is added to the contact point with the tip as **green**, and the (scaled) contact force tip is drawn in **blue**.  Contact points on the bodies are drawn in **orange**."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XLAsUaA1JlGC"
      },
      "source": [
        "shapes = [\"Point\", \"Sphere\", \"Cylinder\", \"Box\", \"Capsule\", \"Ellipsoid\"]\n",
        "shapeA = Dropdown(options=shapes,\n",
        "                  description=\"Body A\",\n",
        "                  value=\"Box\",\n",
        "                  style = {'description_width': 'initial'})\n",
        "shapeB = Dropdown(options=shapes,\n",
        "                  description=\"Body B\",\n",
        "                  style = {'description_width': 'initial'})\n",
        "display(shapeA)\n",
        "display(shapeB)\n",
        "\n",
        "def get_shape(shape):\n",
        "    if shape == \"Point\":\n",
        "        return Sphere(0.01)\n",
        "    elif shape == \"Sphere\":\n",
        "        return Sphere(1.0)\n",
        "    elif shape == \"Cylinder\":\n",
        "        return Cylinder(1.0, 2.0)\n",
        "    elif shape == \"Box\":\n",
        "        return Box(1.0, 2.0, 3.0)\n",
        "    elif shape == \"Capsule\":\n",
        "        return Capsule(1.0, 2.0)\n",
        "    elif shape == \"Ellipsoid\":\n",
        "        return Ellipsoid(1.0, 2.0, 3.0) \n",
        "    raise Exception(\"Unknown shape\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "648WXpNpJlGF"
      },
      "source": [
        "def contact_inspector(shapeA, shapeB):\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    shapeA = get_shape(shapeA.value)\n",
        "    shapeB = get_shape(shapeB.value)\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)\n",
        "    frame = plant.AddFrame(FixedOffsetFrame(\"planar_joint_frame\", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))\n",
        "\n",
        "    mu = 0.5\n",
        "    bodyA = plant.AddRigidBody(\"A\", SpatialInertia(mass=1.0, p_PScm_E=np.array([0., 0., 0.]),\n",
        "            G_SP_E=UnitInertia(1.0, 1.0, 1.0)))\n",
        "    plant.RegisterCollisionGeometry(bodyA, RigidTransform(), shapeA, \"A\", CoulombFriction(mu, mu))\n",
        "    plant.RegisterVisualGeometry(bodyA, RigidTransform(), shapeA, \"A\", [.9, .5, .5, 0.5])\n",
        "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"A\"))\n",
        "\n",
        "    bodyB = plant.AddRigidBody(\"B\", SpatialInertia(mass=1.0, p_PScm_E=np.array([0., 0., 0.]),\n",
        "            G_SP_E=UnitInertia(1.0, 1.0, 1.0)))\n",
        "    plant.RegisterCollisionGeometry(bodyB, RigidTransform(), shapeB, \"B\", CoulombFriction(mu, mu))\n",
        "    plant.RegisterVisualGeometry(bodyB, RigidTransform(), shapeB, \"B\", [.5, .5, .9, 0.9])\n",
        "    plant.AddJoint(PlanarJoint(\"B\", frame, plant.GetFrameByName(\"B\"), damping=[0,0,0]))\n",
        "\n",
        "    plant.Finalize()\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(\n",
        "        builder, \n",
        "        scene_graph, \n",
        "        zmq_url=zmq_url_2d,\n",
        "        server_args=server_args)\n",
        "    meshcat.set_planar_viewpoint(xmin=-3.0, xmax=3.0, ymin=-.2, ymax=3.0)\n",
        "\n",
        "#    contact_visualizer = builder.AddSystem(MeshcatContactVisualizer(meshcat, \n",
        "#                                           force_threshold=0.0,\n",
        "#                                           contact_force_scale=1.0, \n",
        "#                                           contact_force_radius=0.002,   \n",
        "#                                           plant=plant))\n",
        "#    builder.Connect(scene_graph.get_pose_bundle_output_port(),\n",
        "#                    contact_visualizer.GetInputPort(\"pose_bundle\"))\n",
        "#    builder.Connect(plant.get_contact_results_output_port(),\n",
        "#                    contact_visualizer.GetInputPort(\"contact_results\"))\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    meshcat.load()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "\n",
        "    contact = meshcat.vis[\"contact\"]\n",
        "    red = g.MeshLambertMaterial(color=0xff0000)\n",
        "    green = g.MeshLambertMaterial(color=0x00ff00)\n",
        "    blue = g.MeshLambertMaterial(color=0x0000ff)\n",
        "    orange = g.MeshLambertMaterial(color=0xffa500)    \n",
        "\n",
        "    textoutput = Textarea(value=\"\", description=\"contact info: \", layout={'width':'800px','height':'100px'}, style={'description_width':'initial'})\n",
        "    display(textoutput)\n",
        "    def my_callback(plant_context):\n",
        "        results = plant.get_contact_results_output_port().Eval(plant_context)\n",
        "        contact.delete()\n",
        "        textoutput.value = \"\" if results.num_point_pair_contacts()>0 else \"no contact\"\n",
        "        for i in range(results.num_point_pair_contacts()):\n",
        "            info = results.point_pair_contact_info(i)\n",
        "            pair = info.point_pair()\n",
        "            contact[str(i)].set_object(g.Sphere(0.02), red)\n",
        "            contact[str(i)].set_transform(tf.translation_matrix(info.contact_point()))\n",
        "            contact[str(i)+\"A\"].set_object(g.Sphere(0.01), orange)\n",
        "            contact[str(i)+\"A\"].set_transform(tf.translation_matrix(pair.p_WCa))\n",
        "            contact[str(i)+\"B\"].set_object(g.Sphere(0.01), orange)\n",
        "            contact[str(i)+\"B\"].set_transform(tf.translation_matrix(pair.p_WCb))\n",
        "            contact[str(i)+\"normal\"].set_object(g.Sphere(0.02), green)\n",
        "            contact[str(i)+\"normal\"].set_transform(tf.translation_matrix(info.contact_point()-pair.nhat_BA_W ))\n",
        "            contact[str(i)+\"force\"].set_object(g.Sphere(0.02), blue)\n",
        "            contact[str(i)+\"force\"].set_transform(tf.translation_matrix(info.contact_point()+info.contact_force()/5000.0 ))\n",
        "            textoutput.value += f\"slip speed:{info.slip_speed():.4f}, separation_speed:{info.separation_speed():.4f}, depth:{pair.depth:.4f},\\n\"\n",
        "            textoutput.value += \"  point:\" + np.array2string(info.contact_point(), formatter={'float': lambda x: \"{:5.2f}\".format(x)}) + \"\\n\"\n",
        "            textoutput.value += \"  normal:\" + np.array2string(-pair.nhat_BA_W, formatter={'float': lambda x: \"{:5.2f}\".format(x)}) + \"\\n\"\n",
        "            textoutput.value += \"  force:\" + np.array2string(info.contact_force(), formatter={'float': lambda x: \"{:5.2f}\".format(x)}) + \"\\n\"\n",
        "\n",
        "\n",
        "    lower_limit = [-3, -3, -np.pi/2.0]\n",
        "    upper_limit = [3, 3, np.pi/2.0]\n",
        "#    lower_limit += lower_limit\n",
        "#    upper_limit += upper_limit\n",
        "    plant.SetPositions(\n",
        "        plant.GetMyContextFromRoot(context),\n",
        "#        [0, 0, 0] +\n",
        "        [1.2, 1.2, 0.0],\n",
        "    )\n",
        "\n",
        "    MakeJointSlidersThatPublishOnCallback(\n",
        "        plant, diagram, context, my_callback, \n",
        "        lower_limit=lower_limit, upper_limit=upper_limit)    \n",
        "\n",
        "contact_inspector(shapeA, shapeB);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQYb5WgpJlGJ"
      },
      "source": [
        "# Contact geometry of the foam brick"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "mdD-VdxlJlGJ"
      },
      "source": [
        "def show_brick_contact():\n",
        "    builder = DiagramBuilder()\n",
        "\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.01)\n",
        "    Parser(plant).AddModelFromFile(FindResource(\"models/061_foam_brick.sdf\"))\n",
        "    plant.Finalize()\n",
        "\n",
        "    meshcat = ConnectMeshcatVisualizer(\n",
        "        builder, \n",
        "        scene_graph, \n",
        "        zmq_url=zmq_url,\n",
        "        server_args=server_args)\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    meshcat.load()\n",
        "\n",
        "show_brick_contact()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "source": [
        "# Point cloud processing with Open3d\n",
        "\n",
        "I've produced a scene with multiple cameras looking at our favorite YCB mustard bottle.  I've taken the individual point clouds, converted them into Open3d's point cloud format, estimated their normals, merged the point clouds, cropped then point clouds (to get rid of the geometry from the other cameras), then downsampled the point clouds.  (The order is important!)\n",
        "\n",
        "I've pushed all of the point clouds to meshcat, but with many of them set to not be visible by default.  Use the drop-down menu to turn them on and off, and make sure you understand basically what is happening on each of the steps."
      ],
      "cell_type": "markdown",
      "metadata": {
        "id": "qxboQZ1IJlGM"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "def point_cloud_processing_example():\n",
        "    # This just sets up our mustard bottle with three depth cameras positioned around it.\n",
        "    system = MustardExampleSystem()\n",
        "\n",
        "    plant = system.GetSubsystemByName(\"plant\")\n",
        "\n",
        "    # Evaluate the camera output ports to get the images.\n",
        "    context = system.CreateDefaultContext()\n",
        "    plant_context = plant.GetMyContextFromRoot(context)\n",
        "\n",
        "    v = meshcat.Visualizer(zmq_url=zmq_url)\n",
        "    v.delete()\n",
        "    v[\"/Background\"].set_property(\"visible\", False)\n",
        "\n",
        "    pcd = []\n",
        "    for i in range(3):\n",
        "        point_cloud = system.GetOutputPort(f\"camera{i}_point_cloud\").Eval(context)\n",
        "        cloud = create_open3d_point_cloud(point_cloud)\n",
        "\n",
        "        draw_open3d_point_cloud(v[f\"pointcloud{i}\"], cloud)\n",
        "        v[f\"pointcloud{i}\"].set_property(\"visible\", False)\n",
        "\n",
        "        # Crop to region of interest.\n",
        "        pcd.append(cloud.crop(\n",
        "            o3d.geometry.AxisAlignedBoundingBox(min_bound=[-.3, -.3, -.3],\n",
        "                                                max_bound=[.3, .3, .3])))\n",
        "        draw_open3d_point_cloud(v[f\"pointcloud{i}_cropped\"], pcd[i])\n",
        "        v[f\"pointcloud{i}_cropped\"].set_property(\"visible\", False)\n",
        "\n",
        "        pcd[i].estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(\n",
        "            radius=0.1, max_nn=30))\n",
        "\n",
        "        camera = plant.GetModelInstanceByName(f\"camera{i}\")\n",
        "        body = plant.GetBodyByName(\"base\", camera)\n",
        "        X_C = plant.EvalBodyPoseInWorld(plant_context, body)\n",
        "        pcd[i].orient_normals_towards_camera_location(X_C.translation())\n",
        "\n",
        "    # Merge point clouds.  (Note: You might need something more clever here for\n",
        "    # noisier point clouds; but this can often work!)\n",
        "    merged_pcd = pcd[0] + pcd[1] + pcd[2]\n",
        "    draw_open3d_point_cloud(v[\"merged\"], merged_pcd)\n",
        "\n",
        "    # Voxelize down-sample.  (Note that the normls still look reasonable)\n",
        "    down_sampled_pcd = merged_pcd.voxel_down_sample(voxel_size=0.005)\n",
        "    draw_open3d_point_cloud(v[\"down_sampled\"], down_sampled_pcd, normals_scale=0.01)\n",
        "    # Let the normals be drawn, only turn off the object...\n",
        "    v[\"down_sampled\"][\"<object>\"].set_property(\"visible\", False)\n",
        "\n",
        "    # If we wanted to show it in the open3d visualizer, we would use...\n",
        "    # print(\"Use 'n' to show normals, and '+/-' to change their size.\")\n",
        "    # o3d.visualization.draw_geometries([down_sampled_pcd])\n",
        "\n",
        "\n",
        "point_cloud_processing_example()"
      ]
    },
    {
      "source": [
        "# Estimating normals (and local curvature)\n",
        "\n",
        "TODO: Add the version from depth images (nearest pixels instead of nearest neighbors), and implement it in DepthImageToPointCloud."
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "from ipywidgets import IntSlider\n",
        "\n",
        "def normal_estimation():\n",
        "    system = MustardExampleSystem()\n",
        "    context = system.CreateDefaultContext()\n",
        "\n",
        "    v = meshcat.Visualizer(zmq_url=zmq_url)\n",
        "    v.delete()\n",
        "    v[\"/Background\"].set_property(\"visible\", False)\n",
        "\n",
        "    point_cloud = system.GetOutputPort(\"camera0_point_cloud\").Eval(context)\n",
        "    pcd = create_open3d_point_cloud(point_cloud)\n",
        "    draw_open3d_point_cloud(v[\"point_cloud\"], pcd)\n",
        "\n",
        "    kdtree = o3d.geometry.KDTreeFlann(pcd)\n",
        "    pts = np.asarray(pcd.points)\n",
        "\n",
        "    def my_callback(change):\n",
        "        query = pts[change.new,:]\n",
        "        draw_points(v[\"query\"], query, [0.0, 1.0, 0.0], size=0.002)\n",
        "        (num, indices, distances) = kdtree.search_hybrid_vector_3d(\n",
        "            query=query, radius=0.1, max_nn=40)\n",
        "\n",
        "        neighbors = pts[indices, :]\n",
        "        draw_points(v[\"neighbors\"], neighbors.T, [0.0, 0.0, 1.0], size=0.001)\n",
        "\n",
        "        pstar = np.mean(neighbors,axis=0)\n",
        "        prel = neighbors - pstar\n",
        "        W = np.matmul(prel.T, prel)\n",
        "        w, V = np.linalg.eigh(W)\n",
        "        v[\"least_squares_basis\"].set_object(g.triad(0.01))\n",
        "        v[\"least_squares_basis\"].set_transform(RigidTransform(RotationMatrix(np.fliplr(V)), query).GetAsMatrix4())\n",
        "\n",
        "    slider = IntSlider(min=0, max=pts.shape[0])\n",
        "    slider.observe(my_callback, names='value')\n",
        "    slider.value=4165  # force callback immediately \n",
        "    display(slider)\n",
        "\n",
        "normal_estimation()"
      ]
    },
    {
      "source": [
        "# Scoring grasp candidates"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "def grasp_candidate_cost(plant_context, cloud, plant, scene_graph, scene_graph_context, adjust_X_G=False, textbox=None, meshcat=None):\n",
        "    body = plant.GetBodyByName(\"body\")\n",
        "    X_G = plant.GetFreeBodyPose(plant_context, body)\n",
        "\n",
        "    # Transform cloud into gripper frame\n",
        "    X_GW = X_G.inverse()\n",
        "    pts = np.asarray(cloud.points).T\n",
        "    p_GC = X_GW.multiply(pts)\n",
        "\n",
        "    # Crop to a region inside of the finger box.\n",
        "    crop_min = [-.05, 0.1, -0.00625]\n",
        "    crop_max = [.05, 0.1125, 0.00625]\n",
        "    indices = np.all((crop_min[0] <= p_GC[0,:], p_GC[0,:] <= crop_max[0],\n",
        "                      crop_min[1] <= p_GC[1,:], p_GC[1,:] <= crop_max[1],\n",
        "                      crop_min[2] <= p_GC[2,:], p_GC[2,:] <= crop_max[2]), \n",
        "                     axis=0)\n",
        "\n",
        "    if meshcat:\n",
        "        draw_points(meshcat[\"points\"], pts[:, indices], [1., 0, 0], size=0.01)\n",
        "\n",
        "    if adjust_X_G and np.sum(indices)>0:\n",
        "        p_GC_x = p_GC[0, indices]\n",
        "        p_Gcenter_x = (p_GC_x.min() + p_GC_x.max())/2.0\n",
        "        X_G.set_translation(X_G.translation() + X_G.rotation().multiply([p_Gcenter_x, 0, 0]))\n",
        "        plant.SetFreeBodyPose(plant_context, body, X_G)\n",
        "        X_GW = X_G.inverse()\n",
        "\n",
        "    query_object = scene_graph.get_query_output_port().Eval(scene_graph_context)\n",
        "    # Check collisions between the gripper and the sink\n",
        "    if query_object.HasCollisions():\n",
        "        cost = np.inf\n",
        "        if textbox:\n",
        "            textbox.value = \"Gripper is colliding with the sink!\\n\"\n",
        "            textbox.value += f\"cost: {cost}\"\n",
        "        return cost\n",
        "\n",
        "    # Check collisions between the gripper and the point cloud\n",
        "    margin = 0.0  # must be smaller than the margin used in the point cloud preprocessing.\n",
        "    for pt in cloud.points:\n",
        "        distances = query_object.ComputeSignedDistanceToPoint(pt, threshold=margin)\n",
        "        if distances:\n",
        "            cost = np.inf\n",
        "            if textbox:\n",
        "                textbox.value = \"Gripper is colliding with the point cloud!\\n\"\n",
        "                textbox.value += f\"cost: {cost}\"\n",
        "            return cost\n",
        "\n",
        "\n",
        "    n_GC = X_GW.rotation().multiply(np.asarray(cloud.normals)[indices,:].T)\n",
        "\n",
        "    # Penalize deviation of the gripper from vertical.\n",
        "    # weight * -dot([0, 0, -1], R_G * [0, 1, 0]) = weight * R_G[2,1]\n",
        "    cost = 20.0*X_G.rotation().matrix()[2, 1]\n",
        "\n",
        "    # Reward sum |dot product of normals with gripper x|^2\n",
        "    cost -= np.sum(n_GC[0,:]**2)\n",
        "\n",
        "    if textbox:\n",
        "        textbox.value = f\"cost: {cost}\\n\"\n",
        "        textbox.value += \"normal terms:\" + str(n_GC[0,:]**2)\n",
        "    return cost\n",
        "\n",
        "\n",
        "def process_point_cloud(diagram, context, cameras, bin_name):\n",
        "    \"\"\"A \"no frills\" version of the example above, that returns the down-sampled point cloud\"\"\"\n",
        "    plant = diagram.GetSubsystemByName(\"plant\")\n",
        "    plant_context = plant.GetMyContextFromRoot(context)\n",
        "\n",
        "    # Compute crop box.\n",
        "    bin_instance = plant.GetModelInstanceByName(bin_name)\n",
        "    bin_body = plant.GetBodyByName(\"bin_base\", bin_instance)\n",
        "    X_B = plant.EvalBodyPoseInWorld(plant_context, bin_body)\n",
        "    margin = 0.001  # only because simulation is perfect!\n",
        "    a = X_B.multiply([-.22+0.025+margin, -.29+0.025+margin, 0.015+margin])\n",
        "    b = X_B.multiply([.22-0.1-margin, .29-0.025-margin, 2.0])\n",
        "    crop_min = np.minimum(a,b)\n",
        "    crop_max = np.maximum(a,b)\n",
        "\n",
        "    # Evaluate the camera output ports to get the images.\n",
        "    merged_pcd = o3d.geometry.PointCloud()\n",
        "    for c in cameras:\n",
        "        point_cloud = diagram.GetOutputPort(f\"{c}_point_cloud\").Eval(context)\n",
        "        pcd = create_open3d_point_cloud(point_cloud)\n",
        "\n",
        "        # Crop to region of interest.\n",
        "        pcd = pcd.crop(\n",
        "            o3d.geometry.AxisAlignedBoundingBox(min_bound=crop_min,\n",
        "                                                max_bound=crop_max))    \n",
        "\n",
        "        pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(\n",
        "            radius=0.1, max_nn=30))\n",
        "\n",
        "        camera = plant.GetModelInstanceByName(c)\n",
        "        body = plant.GetBodyByName(\"base\", camera)\n",
        "        X_C = plant.EvalBodyPoseInWorld(plant_context, body)\n",
        "        pcd.orient_normals_towards_camera_location(X_C.translation())\n",
        "        \n",
        "        # Merge point clouds.\n",
        "        merged_pcd += pcd\n",
        "\n",
        "    # Voxelize down-sample.  (Note that the normals still look reasonable)\n",
        "    return merged_pcd.voxel_down_sample(voxel_size=0.005)\n",
        "\n",
        "def make_environment_model(directive=None, draw=False, rng=None, num_ycb_objects=0, bin_name=\"bin0\"):\n",
        "    # Make one model of the environment, but the robot only gets to see the sensor outputs.\n",
        "    if not directive:\n",
        "        directive = FindResource(\"models/two_bins_w_cameras.yaml\")\n",
        "\n",
        "    builder = DiagramBuilder()\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0005)\n",
        "    parser = Parser(plant)\n",
        "    parser.package_map().PopulateFromFolder(FindResource(\"\"))\n",
        "    ProcessModelDirectives(LoadModelDirectives(directive), plant, parser)\n",
        "\n",
        "    for i in range(num_ycb_objects):\n",
        "        object_num = rng.integers(len(ycb))\n",
        "        sdf = FindResourceOrThrow(\"drake/manipulation/models/ycb/sdf/\" + ycb[object_num][1])\n",
        "        parser.AddModelFromFile(sdf, f\"object{i}\")\n",
        "\n",
        "    plant.Finalize()\n",
        "    AddRgbdSensors(builder, plant, scene_graph)\n",
        "\n",
        "    if draw:\n",
        "        meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url, prefix=\"environment\")\n",
        "\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "\n",
        "    if num_ycb_objects > 0:\n",
        "        generator = RandomGenerator(rng.integers(1000))  # this is for c++\n",
        "        plant_context = plant.GetMyContextFromRoot(context)\n",
        "        bin_instance = plant.GetModelInstanceByName(bin_name)\n",
        "        bin_body = plant.GetBodyByName(\"bin_base\", bin_instance)\n",
        "        X_B = plant.EvalBodyPoseInWorld(plant_context, bin_body)\n",
        "        z = 0.1\n",
        "        for body_index in plant.GetFloatingBaseBodies():\n",
        "            tf = RigidTransform(\n",
        "                    UniformlyRandomRotationMatrix(generator),  \n",
        "                    [rng.uniform(-.15,.15), rng.uniform(-.2, .2), z])\n",
        "            plant.SetFreeBodyPose(plant_context, \n",
        "                                plant.get_body(body_index),\n",
        "                                X_B.multiply(tf))\n",
        "            z += 0.1\n",
        "\n",
        "        simulator = Simulator(diagram, context)\n",
        "        simulator.AdvanceTo(1.0 if running_as_notebook else 0.1)\n",
        "    elif draw:\n",
        "        meshcat.load()\n",
        "        diagram.Publish(context)\n",
        "\n",
        "\n",
        "    return diagram, context\n",
        "\n",
        "def grasp_score_inspector():\n",
        "    v = meshcat.Visualizer(zmq_url=zmq_url)\n",
        "    v.delete()\n",
        "    environment, environment_context = make_environment_model(directive = FindResource(\"models/clutter_mustard.yaml\"))\n",
        "\n",
        "    # Another diagram for the objects the robot \"knows about\": gripper, cameras, bins.  Think of this as the model in the robot's head.\n",
        "    builder = DiagramBuilder()\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "    parser = Parser(plant)\n",
        "    parser.package_map().PopulateFromFolder(FindResource(\"\"))\n",
        "    ProcessModelDirectives(LoadModelDirectives(FindResource(\"models/clutter_planning.yaml\")), plant, parser)\n",
        "    plant.Finalize()\n",
        "    \n",
        "    v = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url, prefix=\"planning\")\n",
        "    v.load()\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "\n",
        "    cloud = process_point_cloud(environment, environment_context, [\"camera0\", \"camera1\", \"camera2\"], \"bin0\")\n",
        "    draw_open3d_point_cloud(v.vis[\"cloud\"], cloud, size=0.003)\n",
        "\n",
        "    textbox = Textarea(description=\"cost info\", layout={'width': '800px', 'height': '100px'})\n",
        "    display(textbox)\n",
        "\n",
        "    lower_limit = [-np.pi, -np.pi/4., -np.pi/4., -1, -1, 0]\n",
        "    upper_limit = [0, np.pi/4., np.pi/4., 1, 1, 1]\n",
        "    sliders = MakeJointSlidersThatPublishOnCallback(\n",
        "        plant, diagram, context, my_callback=partial(\n",
        "            grasp_candidate_cost, cloud=cloud, plant=plant, scene_graph=scene_graph,\n",
        "            scene_graph_context=scene_graph.GetMyContextFromRoot(context),\n",
        "            textbox=textbox, meshcat=v.vis[\"cost\"]),\n",
        "        lower_limit=lower_limit, upper_limit=upper_limit)\n",
        "    sliders[0].value = -np.pi/2.0\n",
        "    sliders[3].value = -0.05\n",
        "    sliders[4].value = -.5\n",
        "    sliders[5].value = 0.25\n",
        "\n",
        "grasp_score_inspector()"
      ]
    },
    {
      "source": [
        "# Generating grasp candidates"
      ],
      "cell_type": "markdown",
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "tags": []
      },
      "outputs": [],
      "source": [
        "\n",
        "def generate_grasp_candidate_antipodal(plant_context, cloud, plant, scene_graph, scene_graph_context, rng, meshcat=None):\n",
        "    \"\"\"\n",
        "    Picks a random point in the cloud, and aligns the robot finger with the normal of that pixel. \n",
        "    The rotation around the normal axis is drawn from a uniform distribution over [min_roll, max_roll].\n",
        "    \"\"\"\n",
        "    body = plant.GetBodyByName(\"body\")\n",
        "\n",
        "    index = rng.integers(0,len(cloud.points)-1)\n",
        "\n",
        "    # Use S for sample point/frame.\n",
        "    p_WS = np.asarray(cloud.points[index])\n",
        "    n_WS = np.asarray(cloud.normals[index])\n",
        "\n",
        "    if meshcat:\n",
        "        vertices = np.empty((3,2))\n",
        "        vertices[:, 0] = p_WS\n",
        "        vertices[:, 1] = p_WS + 0.05*n_WS\n",
        "        meshcat.set_object(g.LineSegments(g.PointsGeometry(vertices),\n",
        "                            g.MeshBasicMaterial(color=0xff0000)))\n",
        "\n",
        "    assert np.isclose(np.linalg.norm(n_WS), 1.0)\n",
        "\n",
        "    Gx = n_WS # gripper x axis aligns with normal\n",
        "    # make orthonormal y axis, aligned with world down\n",
        "    y = np.array([0.0, 0.0, -1.0])\n",
        "    if np.abs(np.dot(y,Gx)) < 1e-6:\n",
        "        # normal was pointing straight down.  reject this sample.\n",
        "        return None\n",
        "\n",
        "    Gy = y - np.dot(y,Gx)*Gx\n",
        "    Gz = np.cross(Gx, Gy)\n",
        "    R_WG = RotationMatrix(np.vstack((Gx, Gy, Gz)).T)\n",
        "    p_GS_G = [0.054 - 0.01, 0.10625, 0]\n",
        "\n",
        "    # Try orientations from the center out\n",
        "    min_roll=-np.pi/3.0\n",
        "    max_roll=np.pi/3.0\n",
        "    alpha = np.array([0.5, 0.65, 0.35, 0.8, 0.2, 1.0, 0.0])\n",
        "    for theta in (min_roll + (max_roll - min_roll)*alpha):\n",
        "        # Rotate the object in the hand by a random rotation (around the normal).\n",
        "        R_WG2 = R_WG.multiply(RotationMatrix.MakeXRotation(theta))\n",
        "\n",
        "        # Use G for gripper frame.\n",
        "        p_SG_W = - R_WG2.multiply(p_GS_G)\n",
        "        p_WG = p_WS + p_SG_W \n",
        "\n",
        "        X_G = RigidTransform(R_WG2, p_WG)\n",
        "        plant.SetFreeBodyPose(plant_context, body, X_G)\n",
        "        cost = grasp_candidate_cost(plant_context, cloud, plant, scene_graph, scene_graph_context, adjust_X_G=True, meshcat=meshcat)\n",
        "        X_G = plant.GetFreeBodyPose(plant_context, body)\n",
        "        if np.isfinite(cost):\n",
        "            return cost, X_G\n",
        "\n",
        "        #draw_grasp_candidate(X_G, f\"collision/{theta:.1f}\")\n",
        "\n",
        "    return np.inf, None\n",
        "\n",
        "def draw_grasp_candidate(X_G, prefix='gripper', draw_frames=True):\n",
        "    builder = DiagramBuilder()\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "    parser = Parser(plant)\n",
        "    parser.package_map().Add(\"wsg_50_description\", os.path.dirname(FindResourceOrThrow(\"drake/manipulation/models/wsg_50_description/package.xml\")))\n",
        "    gripper = parser.AddModelFromFile(FindResource(\n",
        "        \"models/schunk_wsg_50_welded_fingers.sdf\"), \"gripper\")\n",
        "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"body\"), X_G)\n",
        "    plant.Finalize()\n",
        "    \n",
        "    frames_to_draw = {\"gripper\": {\"body\"}} if draw_frames else {}\n",
        "    meshcat = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url, prefix=prefix, delete_prefix_on_load=False, frames_to_draw=frames_to_draw)\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "\n",
        "    meshcat.load()\n",
        "    diagram.Publish(context)\n",
        "\n",
        "def sample_grasps_example():\n",
        "    v = meshcat.Visualizer(zmq_url=zmq_url)\n",
        "    v.delete()\n",
        "    rng = np.random.default_rng()\n",
        "\n",
        "    environment, environment_context = make_environment_model(rng=rng, num_ycb_objects=5, draw=False)\n",
        "\n",
        "    # Another diagram for the objects the robot \"knows about\": gripper, cameras, bins.  Think of this as the model in the robot's head.\n",
        "    builder = DiagramBuilder()\n",
        "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
        "    parser = Parser(plant)\n",
        "    parser.package_map().PopulateFromFolder(FindResource(\"\"))\n",
        "    ProcessModelDirectives(LoadModelDirectives(FindResource(\"models/clutter_planning.yaml\")), plant, parser)\n",
        "    plant.Finalize()\n",
        "    \n",
        "    v = ConnectMeshcatVisualizer(builder, scene_graph, zmq_url=zmq_url, prefix=\"planning\")\n",
        "    v.load()\n",
        "    diagram = builder.Build()\n",
        "    context = diagram.CreateDefaultContext()\n",
        "    diagram.Publish(context)\n",
        "    v.vis[\"planning/10\"].set_property('visible', False)  # Hide this particular gripper\n",
        "\n",
        "    cloud = process_point_cloud(environment, environment_context, [\"camera0\", \"camera1\", \"camera2\"], \"bin0\")\n",
        "    draw_open3d_point_cloud(v.vis[\"cloud\"], cloud, size=0.003)\n",
        "\n",
        "    plant_context = plant.GetMyContextFromRoot(context)\n",
        "    scene_graph_context = scene_graph.GetMyContextFromRoot(context)\n",
        "\n",
        "    costs = []\n",
        "    X_Gs = []\n",
        "    for i in range(100):\n",
        "        cost, X_G = generate_grasp_candidate_antipodal(plant_context, cloud, plant, scene_graph, scene_graph_context, rng)#, meshcat=v.vis[\"sample\"])\n",
        "        if np.isfinite(cost):\n",
        "            costs.append(cost)\n",
        "            X_Gs.append(X_G)\n",
        "\n",
        "    indices = np.asarray(costs).argsort()[:5]\n",
        "    for i in indices:\n",
        "        draw_grasp_candidate(X_Gs[i], prefix=f\"{i}th best\", draw_frames=False)\n",
        "\n",
        "\n",
        "sample_grasps_example()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ]
}